from model.province_city import ProvinceCity
from sqlalchemy import asc, and_
from utils.work_db import get_session


class ProvinceCityDao(object):
    def __init__(self, session):
        self.session = session

    def add(self, job):
        self.session.add(job)

    def get_by_id(self, pc_id):
        return self.session.query(ProvinceCity).filter(ProvinceCity.pc_id == pc_id).one()

    """获得顶级城市信息"""

    def list_province(self):
        q = self.session.query(ProvinceCity).filter(ProvinceCity.parent_id == 0)
        q = q.order_by(asc(ProvinceCity.code))
        return q.all()

    def list_citiy(self):
        q = self.session.query(ProvinceCity).filter(and_(ProvinceCity.parent_id != 0, ProvinceCity.level != 0))
        q = q.order_by(asc(ProvinceCity.code))
        return q.all()

    def list_city_dict(self):
        # class_code : [{}]
        fdict = {}
        list = self.list_citiy()
        if list:
            for d in list:
                s_list = []
                if not fdict.get(str(d.parent_id)):
                    s_list.append({"code": d.code, "name": d.name})
                    fdict[d.parent_id] = s_list
                else:
                    s_list = fdict.get(str(d.parent_id))
                    s_list.append({"code": d.code, "name": d.name})

            return fdict
        else:
            return None

    def list_all(self):
        datas = []
        result = self.list_province()
        if result:
            children_dict = self.list_city_dict()
            for d in result:
                datas.append(
                    {"code": d.code,
                     "name": d.name,
                     "children": children_dict.get(str(d.pc_id))})
            return datas
        else:
            return None


if __name__ == '__main__':
    s = get_session()
    dao = ProvinceCityDao(s)
    print(dao.list_all())
